use actix_web::{
    dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform},
    Error, HttpMessage,
};
use futures_util::future::LocalBoxFuture;
use std::{
    future::{ready, Ready},
    rc::Rc,
};
use uuid::Uuid;

pub struct RequestId;

impl<S, B> Transform<S, ServiceRequest> for RequestId
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type InitError = ();
    type Transform = RequestIdService<S>;
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ready(Ok(RequestIdService {
            service: Rc::new(service),
        }))
    }
}

pub struct RequestIdService<S> {
    service: Rc<S>,
}

impl<S, B> Service<ServiceRequest> for RequestIdService<S>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
    S::Future: 'static,
    B: 'static,
{
    type Response = ServiceResponse<B>;
    type Error = Error;
    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;

    forward_ready!(service);

    fn call(&self, req: ServiceRequest) -> Self::Future {
        let service = self.service.clone();

        Box::pin(async move {
            // Generate or extract request ID
            let request_id = req
                .headers()
                .get("x-request-id")
                .and_then(|h| h.to_str().ok())
                .and_then(|s| Uuid::parse_str(s).ok())
                .unwrap_or_else(|| Uuid::new_v4());

            // Store request ID in extensions
            req.extensions_mut().insert(request_id);

            // Call the service
            let mut res = service.call(req).await?;

            // Add request ID to response headers
            res.headers_mut().insert(
                actix_web::http::header::HeaderName::from_static("x-request-id"),
                actix_web::http::header::HeaderValue::from_str(&request_id.to_string())
                    .unwrap_or_else(|_| actix_web::http::header::HeaderValue::from_static("invalid")),
            );

            Ok(res)
        })
    }
}

// Helper function to get request ID from request
pub fn get_request_id(req: &ServiceRequest) -> Option<Uuid> {
    req.extensions().get::<Uuid>().copied()
}

// Helper function for handlers to get request ID
pub fn get_request_id_from_http_req(req: &actix_web::HttpRequest) -> Option<Uuid> {
    req.extensions().get::<Uuid>().copied()
}

#[cfg(test)]
mod tests {
    use super::*;
    use actix_web::{test, web, App, HttpResponse};

    async fn test_handler() -> HttpResponse {
        HttpResponse::Ok().json(serde_json::json!({"message": "test"}))
    }

    #[actix_web::test]
    async fn test_request_id_middleware() {
        let app = test::init_service(
            App::new()
                .wrap(RequestId)
                .route("/test", web::get().to(test_handler))
        ).await;

        let req = test::TestRequest::get()
            .uri("/test")
            .to_request();
            
        let resp = test::call_service(&app, req).await;
        assert!(resp.status().is_success());
        
        // Check if response has request ID header
        let headers = resp.headers();
        assert!(headers.contains_key("x-request-id"));
    }

    #[actix_web::test]
    async fn test_request_id_middleware_with_existing_id() {
        let app = test::init_service(
            App::new()
                .wrap(RequestId)
                .route("/test", web::get().to(test_handler))
        ).await;

        let existing_id = Uuid::new_v4();
        let req = test::TestRequest::get()
            .uri("/test")
            .insert_header(("x-request-id", existing_id.to_string()))
            .to_request();
            
        let resp = test::call_service(&app, req).await;
        assert!(resp.status().is_success());
        
        // Check if response has the same request ID
        let headers = resp.headers();
        let response_id = headers.get("x-request-id").unwrap().to_str().unwrap();
        assert_eq!(response_id, existing_id.to_string());
    }
}